#ifndef YOLOV5_H
#define YOLOV5_H

#include <iostream>
#include <string>
#include <opencv2/opencv.hpp>
#include <unistd.h>
#include <omp.h>
#include<chrono>

#include "rt_model_operation.h"
using namespace algocommon;

namespace yolov5
{
    struct InitParams:public OnnxModeltoTRTParam
    {
        int gpu_id;
        std::string onnx_model_path;
        std::string trt_model_dir;
        std::string trt_model_name;
        int longsize = 640;
        std::vector<std::vector<std::vector<int> > >anchor_size = {{{10,13},{16,30}, {33,23}},
                                                                    {{30,61}, {62,45}, {59,119}},
                                                                   {{116,90}, {156,198}, {373,326}}};
    };

    struct Anchor
    {
        float x;
        float y;
        float w;
        float h;
        float score;
        int cls;
    };

    class Yolov5 : public OnnxModeltoTRT
    {
    public:

        Yolov5(const InitParams &params);
        virtual ~Yolov5(){};

    public:

        std::vector<std::vector<float> >  Execute(const cv::Mat &img);
        std::vector<std::vector<float> >  ExecuteGPU(const cv::Mat &img);
        
        cudaStream_t stream_;

    private:

        Result InitNetwork();
        void PreProcessCPU(const cv::Mat &img);
        void PreProcessCUDA(const cv::Mat &img);
        std::vector<Anchor> PostProcessCPU(std::vector<TensorFloat>&output_vectors);
        std::vector<Anchor> PostProcessCUDA(std::vector<TensorFloat>&output_vectors);
        Result NMS(std::vector<Anchor> &anchors);
        std::vector<Anchor> DiouNms(std::vector<Anchor> &anchors);
        // cv::Mat PostOutput(const std::vector<cv::Mat>&output_vector);
        void SetCutSize(const cv::Mat &img);

        int64_t CurrentTimeMillis()
        {
            int64_t timems = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now().time_since_epoch()).count();
            return timems;
        }
    private:

        InitParams params_;

        std::vector<int> stride_ = {8, 16, 32};

        std::vector<std::vector<std::vector<int> > >anchor_size_ = {{{10,13},{16,30}, {33,23}},
                                                                    {{30,61}, {62,45}, {59,119}},
                                                                   {{116,90}, {156,198}, {373,326}}};

        // std::vector<std::vector<std::vector<int> > >anchor_size_ = {{{36,20},{10,114}, {119,13}},
        //                                                             {{26,154}, {18,422}, {159,47}},
        //                                                            {{401,30}, {35,471}, {431,72}}};
        
        std::vector<std::vector<int> > anchor_size_cuda_ = {{10, 13, 16, 30, 33, 23},
                                                           {30, 61, 62, 45, 59, 119},
                                                           {116, 90, 156, 198, 373, 32}};

        TensorUint8 src_cuda_gloablmem_;     // 原始图像大小
        TensorUint8 resize_cuda_gloablmem_;  // 原始图像尺度变换后的大小
        TensorUint8 cut_cuda_gloablmem_;     // CUT图像尺度变换后的大小
    
        

        // int longsize_ = 512;
        int longsize_ = 640;
        // int longsize_ = 480;
        int nms_bbox_thresh_ = 0.3;
        int nms_score_thresh_ = 0.3;
        float score_thresh_ = 0.3;
        int ow_; //cut_image
        int oh_;

    };























}//namespace yolov5

















#endif